package com.xl.bigdata.es.search;

import com.xl.bigdata.es.xpack.ESXpackRestHighLevelClient;
import org.elasticsearch.action.search.*;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.TermQueryBuilder;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;

import java.io.IOException;
import java.util.*;

/**
 * @ClassName SearchScrollAPI
 * @Description TODO
 * 滚动查询大量历史数据
 * 官网文档：https://www.elastic.co/guide/en/elasticsearch/reference/current/paginate-search-results.html
 * 官网高级JavaAPI文档：https://www.elastic.co/guide/en/elasticsearch/client/java-rest/7.15/java-rest-high-search-scroll.html
 * @Author xl.gao
 * @Date 2021/12/28 18:08
 * @Version
 */
public class SearchScrollAPI {

    public static final int SCROLL_TIMEOUT_SECOND = 30;

    public static Map<String, List<String>> searchSrcollAPI(RestHighLevelClient client, String index, int size, String scrollId){
        Map<String, List<String>> map = new HashMap<>();

        if(scrollId == null || "".equals(scrollId)){
            map = searchSrollFirstAPI(client, index, size);
        } else {
            map = searchSrollNextAPI(client, index, size, scrollId);
        }

        return map;

    }

    public static Map<String, List<String>> searchSrcollAPI(RestHighLevelClient client, String index, int size, String scrollId, QueryBuilder qb){
        Map<String, List<String>> map = new HashMap<>();

        if(scrollId == null || "".equals(scrollId)){
            map = searchSrollFirstAPI(client, index, size, qb);
        } else {
            map = searchSrollNextAPI(client, index, size, scrollId);
        }

        return map;

    }

    /*
     * 第二步，必须将接收到的滚动标识符设置为 SearchScrollRequest 以及新的滚动间隔，
     * 并通过 searchScroll 方法发送。 Elasticsearch 返回另一批带有新滚动标识符的结果。
     * 然后可以在后续的 SearchScrollRequest 中使用这个新的滚动标识符来检索下一批结果，
     * 依此类推。这个过程应该循环重复，直到没有更多的结果返回，这意味着滚动已经用完并且所有匹配的文档都已经被检索到。
     */
    public static Map<String, List<String>> searchSrollNextAPI(RestHighLevelClient client, String index, int size, String scrollId){

        Map<String, List<String>> map = new HashMap<>();
        List<String> list = new ArrayList<>();

        SearchScrollRequest searchScrollRequest = new SearchScrollRequest(scrollId);
        searchScrollRequest.scroll(TimeValue.timeValueSeconds(SCROLL_TIMEOUT_SECOND));

        try {
            SearchResponse searchResponse = client.scroll(searchScrollRequest, RequestOptions.DEFAULT);
            String scrollId1 = searchResponse.getScrollId();
            SearchHits hits = searchResponse.getHits();
            hits.forEach( sh -> {
                String sourceAsString = sh.getSourceAsString();
                list.add(sourceAsString);
            });

            map.put(scrollId, list);
        } catch (IOException e) {
            e.printStackTrace();
        }


        return map;
    }

    /*
     * 第一步：必须执行一个带有滚动参数的初始搜索请求，以通过搜索API初始化滚动会话。在处理这个SearchRequest时，
     * Elasticsearch会检测到滚动参数的存在，并在相应的时间间隔内保持搜索上下文。
     */
    public static Map<String, List<String>> searchSrollFirstAPI(RestHighLevelClient client, String index, int size){
        Map<String, List<String>> map = new HashMap<>();
        List<String> list = new ArrayList<>();

        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(QueryBuilders.termQuery("citycomple", "青岛"));
        searchSourceBuilder.size(size);
//        searchSourceBuilder.

        searchRequest.source(searchSourceBuilder);
        searchRequest.scroll(TimeValue.timeValueMinutes(1L));

        try {
            SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);

            String scrollId = searchResponse.getScrollId();

            SearchHits hits = searchResponse.getHits();
            hits.forEach( sh -> {
                String sourceAsString = sh.getSourceAsString();
                list.add(sourceAsString);
            });

            map.put(scrollId, list);

        } catch (IOException e) {
            e.printStackTrace();
        }

        return map;
    }

    /*
     * 第一步：必须执行一个带有滚动参数的初始搜索请求，以通过搜索API初始化滚动会话。在处理这个SearchRequest时，
     * Elasticsearch会检测到滚动参数的存在，并在相应的时间间隔内保持搜索上下文。
     */
    public static Map<String, List<String>> searchSrollFirstAPI(RestHighLevelClient client, String index, int size, QueryBuilder qb){
        Map<String, List<String>> map = new HashMap<>();
        List<String> list = new ArrayList<>();

        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(qb);
        searchSourceBuilder.size(size);
//        searchSourceBuilder.

        searchRequest.source(searchSourceBuilder);
        searchRequest.scroll(TimeValue.timeValueMinutes(1L));

        try {
            SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);

            String scrollId = searchResponse.getScrollId();

            SearchHits hits = searchResponse.getHits();
            hits.forEach( sh -> {
                String sourceAsString = sh.getSourceAsString();
                list.add(sourceAsString);
            });

            map.put(scrollId, list);

        } catch (IOException e) {
            e.printStackTrace();
        }

        return map;
    }

    /*
     * 必须关闭滚动搜索
     * Search Scroll API 使用的搜索上下文会在滚动超时时自动删除。但建议在不再需要时使用 Clear Scroll API 释放搜索上下文。
     */
    public static boolean closeSearchScrollAPI(RestHighLevelClient client, String scrollId){

        ClearScrollRequest clearScrollRequest = new ClearScrollRequest();

        clearScrollRequest.addScrollId(scrollId);

        try {
            ClearScrollResponse clearScrollResponse = client.clearScroll(clearScrollRequest, RequestOptions.DEFAULT);

            boolean succeeded = clearScrollResponse.isSucceeded();

            int numFreed = clearScrollResponse.getNumFreed();

            return succeeded;
        } catch (IOException e) {
            e.printStackTrace();
        }

        return false;
    }

    public static void main(String[] args) {

        TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery("citycomple", "青岛");

        RestHighLevelClient restHighLevelClient = ESXpackRestHighLevelClient.getRestHighLevelClient();

        String scrollId = "";
        for (int i = 0; i < 10; i++) {

            Map<String, List<String>> stringListMap = searchSrcollAPI(restHighLevelClient, "address_base", 10000, scrollId, termQueryBuilder);

            for(String k : stringListMap.keySet()){
                scrollId = k;
            }

            System.out.println("scrollId : " + scrollId + " mapSize : " + stringListMap.get(scrollId).size());
        }

        closeSearchScrollAPI(restHighLevelClient, scrollId);

        ESXpackRestHighLevelClient.closeRestHighLevelClient();
    }
}
